import torch  # 导入 PyTorch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# ✅ 修改模型保存路径，尝试直接从 test_trainer 目录加载
model_path = "../test_trainer"

# ✅ 先检查 tokenizer 是否存在
try:
    trained_tokenizer = AutoTokenizer.from_pretrained(model_path)
except OSError:
    print("❌ 无法加载 tokenizer，请检查路径或重新训练模型")
    exit()

# ✅ 加载训练好的模型
trained_model = AutoModelForSequenceClassification.from_pretrained(model_path)

# ✅ 检测设备并移动模型到 GPU/CPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
trained_model.to(device)  # 确保模型在正确的设备上


def predict_sentiment(text):
    """
    使用训练好的模型预测文本情感。

    Args:
        text: 要预测的文本 (字符串)。

    Returns:
        预测的标签 (整数，0-4，对应 1-5 星评分) 和 概率 (PyTorch 张量)。
    """

    # 1. 使用分词器处理输入文本
    inputs = trained_tokenizer(text, padding=True, truncation=True, return_tensors="pt")

    # ✅ 将输入数据移动到与模型相同的设备上 (CPU 或 GPU)
    inputs = {key: value.to(device) for key, value in inputs.items()}  # 确保 inputs 在正确的设备上

    # 3. 使用模型进行预测 (前向传播)
    with torch.no_grad():
        outputs = trained_model(**inputs)

    # 4. 获取预测的 logits (模型的原始输出)
    logits = outputs.logits

    # 5. 将 logits 转换为概率 (使用 softmax 函数)
    probabilities = torch.softmax(logits, dim=-1)

    # 6. 获取预测的标签 (概率最高的类别)
    predicted_label_index = torch.argmax(logits, dim=-1).item()

    return predicted_label_index, probabilities


# 示例文本
sample_text = "Today is a nice day!"
sample_text_negative = "This place was terrible. The food was cold and the waiter was rude. Never going back."

# ✅ 进行预测
predicted_label, probabilities = predict_sentiment(sample_text)
predicted_label_negative, probabilities_negative = predict_sentiment(sample_text_negative)

# ✅ 标签 (0-4) 对应 1-5 星评分，所以需要 +1
predicted_star_rating = predicted_label + 1
predicted_star_rating_negative = predicted_label_negative + 1

print(f"✅ 正面评价: \"{sample_text}\"")
print(f"预测星级: {predicted_star_rating} 星")
print(f"概率分布: {probabilities}")

print(f"\n✅ 负面评价: \"{sample_text_negative}\"")
print(f"预测星级: {predicted_star_rating_negative} 星")
print(f"概率分布: {probabilities_negative}")
